import os
import torch
from torch import nn
import pandas as pd
from torch.utils.data import DataLoader
from dataset.Data_Loader import LOAD_DATA
from model.ISRG_block import ISRG
from params import params_dict
from tqdm import tqdm
import sys
from utils.evaluate import *
from utils.utils import *


device = torch.device("cpu")
model = load_model(params_dict)
model.load_state_dict(torch.load('run/best5.pt'))
model.to(device)
test_dataset = LOAD_DATA(params_dict['test_img_path'], params_dict['test_file_path'], mode='test1.txt')
# test_dataset = LOAD_DATA(params_dict['valid_img_path'], params_dict['valid_file_path'], mode='test1.txt')
test_num = len(test_dataset)
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
model.eval()
class_list = []
label_list = []
with torch.no_grad():
    test_bar = tqdm(test_loader, file=sys.stdout)
    for test_data in test_bar:
        test_imgs, test_features_seqs, test_labels, _ = test_data
        outputs = model(test_imgs.to(params_dict['device']), test_features_seqs.to(params_dict['device']))
        pred = torch.softmax(outputs, dim=1)
        predict_cla = torch.argmax(pred, dim=1).numpy()
        class_list.extend(predict_cla)
        label_list.extend(list(map(int, test_labels)))

pred_tensor = torch.tensor(class_list)
true_tensor = torch.tensor(label_list)
print("F1-score")
print(test_f1_en(pred_tensor, true_tensor))

print("Accuracy")
print(test_acc_en(pred_tensor, true_tensor))

print("Specificity")
print(test_spc_en(pred_tensor, true_tensor))

print("recall")
print(test_rcl_en(pred_tensor, true_tensor))

print("Precision")
print(test_pcl_en(pred_tensor, true_tensor))

print("Kapa")
print(cohen_kappa_score(pred_tensor, true_tensor))

matrix(class_list, label_list, ["noprolapsed", "mild", "middle", "serious"])

